from leetcode_test import TreeNode

"""
并不是要求所有节点都满足绝对差小于等于1,而是左右有一个满足条件就行
"""


class Solution:
    def isBalanced(self, root: TreeNode) -> bool:
        def helper(root):
            if not root:
                return True, 0

            left, left_depth = helper(root.left)
            right, right_depth = helper(root.right)
            if left and right and abs(left_depth - right_depth) <= 1:
                return True, max(left_depth, right_depth) + 1
            else:
                return False, 0

        return helper(root)[0]


# class Solution(object):
#     def isBalanced(self, root):

#         def isB(root):
#             if not root:return True,0
#             lb,ls = isB(root.left)
#             rb,rs = isB(root.right)
#             print(lb,ls,rb,rs)
#             return lb and rb and abs(ls - rs) <= 1, max(ls, rs) + 1

#         return isB(root)[0]


def test():
    t = TreeNode.create([1, 2, 2, 3, 3, None, None, 4, 4])
    assert Solution().isBalanced(t) == False


def test1():
    t = TreeNode.create([3, 9, 20, None, None, 15, 7])
    assert Solution().isBalanced(t)


def test2():
    t = TreeNode.create([1, 2, 2, 3, None, None, 3, 4, None, None, 4])
    # 节点2为例, 左边的长度是2,右面节点的长度是0,所以不是
    assert Solution().isBalanced(t) == False


def test3():
    t = TreeNode.create([1, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, None, None, 5, 5])
    assert Solution().isBalanced(t)


if __name__ == "__main__":
    # t = TreeNode.create([3, 9, 20, None, None, 15, 7,3])
    t = TreeNode.create([1, 2, 2, 3, 3, None, None, 4, 4])
    print(Solution().isBalanced(t))
